/**
 * Copyright (C), 2020-2021 by colley.ma
 * File Name: MdcTaskDecorator.java
 * Encoding: UTF-8
 * Date: 2021年8月16日 下午1:35:36
 * History:
 */
package com.swak.frame.monitor.async;

import java.util.Map;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import com.swak.frame.context.ApiContext;
import com.swak.frame.context.ApiContextHolder;
import com.swak.frame.context.ThreadLocalClearUp;

public class ContextCopyingDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        try {
            Map<String, String> previous = MDC.getCopyOfContextMap();
            ApiContext apiContext = ApiContextHolder.get();
            return () -> {
                try {
                    if (previous != null) {
                        MDC.setContextMap(previous);
                    }
                    if (apiContext != null) {
                        ApiContextHolder.set(apiContext);
                    }

                    runnable.run();
                } finally {
                    ThreadLocalClearUp.clear();
                }
            };
        } catch (IllegalStateException e) {
            return runnable;
        }
    }
}
